import torch
import torch.nn as nn
from models.ShangJiaoOfficial.model_mj_in3sadfr import UNet3D as Base
from utils.Utils import print_GPU_memory_usage


class UNet3D(Base):
	"""
	Baseline model with Feature Recalibration module
	for pulmonary artery-vein segmentation
	with deep supervision
	"""
	def __init__(self, in_channels=1, out_channels=3, inputshape=[64, 176, 176], coord=True):
		super(UNet3D, self).__init__(in_channels=in_channels, out_channels=out_channels,\
		                             inputshape=inputshape, coord=coord)
		self.upsampling4 = nn.Upsample(scale_factor=4)
		self.upsampling8 = nn.Upsample(scale_factor=8)

		self.dsconv6 = nn.Conv3d(256, 3, 1, 1, 0)  # deep supervision
		self.dsconv7 = nn.Conv3d(128, 3, 1, 1, 0)  # deep supervision
		self.dsconv8 = nn.Conv3d(64, 3, 1, 1, 0)  # deep supervision

	def forward(self, x, coordmap=None):
		"""
		:param x: shape = (batch_size, num_channels, D, H, W) \
		:param coordmap: shape = (batch_size, num_channels, D, H, W)
		:return: output segmentation tensors list, attention mapping
		"""
		conv1 = self.conv1(x)
		print_GPU_memory_usage("conv1")
		conv1, _ = self.pe1(conv1)  # the conv1 is the output_tensor
		print_GPU_memory_usage("pe1")
		x = self.pooling(conv1)  # kernel_size is 2*2*2
		print_GPU_memory_usage("pooling1")
		
		conv2 = self.conv2(x)
		print_GPU_memory_usage("conv2")
		conv2, _ = self.pe2(conv2)
		print_GPU_memory_usage("pe2")
		x = self.pooling(conv2)
		
		conv3 = self.conv3(x)
		print_GPU_memory_usage("conv3")
		conv3, mapping3 = self.pe3(conv3)
		print_GPU_memory_usage("pe3")
		x = self.pooling(conv3)
		
		conv4 = self.conv4(x)
		print_GPU_memory_usage("conv4")
		conv4, mapping4 = self.pe4(conv4)
		print_GPU_memory_usage("pe4")
		x = self.pooling(conv4)

		conv5 = self.conv5(x)
		print_GPU_memory_usage("conv5")
		conv5, mapping5 = self.pe5(conv5)
		print_GPU_memory_usage("pe5")

		x = self.upsampling(conv5)
		print_GPU_memory_usage("upsampling")
		x = torch.cat([x, conv4], dim=1)
		conv6 = self.conv6(x)
		print_GPU_memory_usage("conv6")
		conv6, mapping6 = self.pe6(conv6)
		print_GPU_memory_usage("pe6")
		
		x = self.upsampling(conv6)
		x = torch.cat([x, conv3], dim=1)
		conv7 = self.conv7(x)
		print_GPU_memory_usage("conv7")
		conv7, mapping7 = self.pe7(conv7)
		print_GPU_memory_usage("pe7")
		
		x = self.upsampling(conv7)
		x = torch.cat([x, conv2], dim=1)
		conv8 = self.conv8(x)
		print_GPU_memory_usage("conv8")
		conv8, mapping8 = self.pe8(conv8)
		print_GPU_memory_usage("pe8")
		
		x = self.upsampling(conv8)
		if self._coord and coordmap is not None:
			x = torch.cat([x, conv1, coordmap], dim=1)
		else:
			x = torch.cat([x, conv1], dim=1)

		conv9 = self.conv9(x)
		print_GPU_memory_usage("conv9")
		conv9, mapping9 = self.pe9(conv9)
		print_GPU_memory_usage("pe9")

		x = self.conv10(conv9)
		print_GPU_memory_usage("conv10")

		ds_6 = self.upsampling8(self.dsconv6(conv6))
		ds_7 = self.upsampling4(self.dsconv7(conv7))
		ds_8 = self.upsampling(self.dsconv8(conv8))

		x2 = self.sigmoid(self.conv11(x))

		return [x, ds_6, ds_7, ds_8], x2, [mapping3, mapping4, mapping5, mapping6, mapping7, mapping8, mapping9]


if __name__ == '__main__':
	net = UNet3D(in_channels=3, out_channels=3, inputshape=[64, 176, 176])
	print(net)
	print('Number of network parameters:', sum(param.numel() for param in net.parameters()))
# Number of network parameters: 16912221

